package websocket

import (
	"common/global"
	"context"
	"encoding/json"
	"equipment_server/model"
	"log"
	"sync"
	"time"
)

// Hub maintains the set of active clients and broadcasts messages to the
// clients.
type Hub struct {
	// Registered clients.
	clients map[uint]*Client

	// Inbound messages from the clients.
	broadcast chan []byte

	// Register requests from the clients.
	register chan *Client

	// Unregister requests from clients.
	unregister chan *Client

	// Mutex to protect clients map
	mu sync.RWMutex
}

var GlobalHub *Hub

func NewHub() *Hub {
	return &Hub{
		broadcast:  make(chan []byte),
		register:   make(chan *Client),
		unregister: make(chan *Client),
		clients:    make(map[uint]*Client),
	}
}

func init() {
	GlobalHub = NewHub()
	go GlobalHub.Run()
}

func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.mu.Lock()
			h.clients[client.UserID] = client
			h.mu.Unlock()
			log.Printf("用户 %d 上线", client.UserID)
		case client := <-h.unregister:
			h.mu.Lock()
			if _, ok := h.clients[client.UserID]; ok {
				delete(h.clients, client.UserID)
				close(client.send)
			}
			h.mu.Unlock()
			log.Printf("用户 %d 下线", client.UserID)
		case message := <-h.broadcast:
			h.mu.RLock()
			for _, client := range h.clients {
				select {
				case client.send <- message:
				default:
					close(client.send)
					delete(h.clients, client.UserID)
				}
			}
			h.mu.RUnlock()
		}
	}
}

// SendToUser sends a message to a specific user.
func (h *Hub) SendToUser(userID uint, msgType string, success bool, payload interface{}) {
	h.mu.RLock()
	client, ok := h.clients[userID]
	h.mu.RUnlock()

	wsMessage := model.WsMessage{
		Type:    msgType,
		Success: success,
		Payload: payload,
	}

	message, err := json.Marshal(wsMessage)
	if err != nil {
		log.Printf("无法序列化 websocket 消息: %v", err)
		return
	}

	if ok {
		// User is online
		select {
		case client.send <- message:
			log.Printf("成功发送消息给用户 %d", userID)
		default:
			// Buffer is full, handle as offline
			log.Printf("用户 %d 的发送缓冲区已满，视为离线处理", userID)
			h.storeOfflineMessage(userID, message)
		}
	} else {
		// User is offline, store the message
		log.Printf("用户 %d 不在线，存储离线消息", userID)
		h.storeOfflineMessage(userID, message)
	}
}

// storeOfflineMessage stores a message for an offline user.
func (h *Hub) storeOfflineMessage(userID uint, message []byte) {
	collection := global.MonGoDatabase.Collection("offline_ws_messages")
	_, err := collection.InsertOne(context.TODO(), model.OfflineWsMessage{
		UserID:    userID,
		Message:   message,
		Timestamp: time.Now().Unix(),
	})
	if err != nil {
		log.Printf("Failed to store offline message for user %d: %v", userID, err)
	} else {
		log.Printf("Stored offline message for user %d", userID)
	}
}
